import { BaseAgent, BaseAgentConfig, AgentCommandConfig } from "./base";
import { ModelConfig } from "./utils";
import {
  GeminiConfig,
  GeminiResponse,
  GeminiStreamCallbacks,
  Conversation,
  ModelProvider,
  AgentType,
} from "../types";

export class GeminiAgent extends BaseAgent {
  private geminiApiKey: string;
  private model?: string;

  private escapePrompt(prompt: string): string {
    // Comprehensive escaping for bash double quotes - production ready
    return prompt
      .replace(/\\/g, '\\\\')    // Escape backslashes FIRST
      .replace(/"/g, '\\"')      // Escape double quotes
      .replace(/'/g, "'\\''")    // Escape single quotes (close quote, escaped quote, reopen)
      .replace(/\$/g, '\\$')     // Escape dollar signs
      .replace(/`/g, '\\`')      // Escape backticks (command substitution)
      .replace(/!/g, '\\!')      // Escape exclamation (history expansion)
      .replace(/\n/g, '\\n')     // Escape newlines
      .replace(/\r/g, '\\r')     // Escape carriage returns
      .replace(/\t/g, '\\t');    // Escape tabs
  }

  constructor(config: GeminiConfig) {
    const baseConfig: BaseAgentConfig = {
      sandboxProvider: config.sandboxProvider,
      secrets: config.secrets,
      sandboxId: config.sandboxId,
      workingDirectory: config.workingDirectory,
    };

    super(baseConfig);

    // Validate that provider is gemini or google if specified
    if (config.provider && !["gemini", "google"].includes(config.provider)) {
      throw new Error(
        "Gemini agent only supports 'gemini' or 'google' provider"
      );
    }

    this.geminiApiKey = config.providerApiKey || "";
    this.model = config.model;
  }

  protected getCommandConfig(
    prompt: string,
    mode?: "ask" | "code"
  ): AgentCommandConfig {
    let instruction: string;
    if (mode === "ask") {
      instruction =
        "Research the repository and answer the user's questions. " +
        "Do NOT make any changes to any files in the repository.";
    } else {
      instruction =
        "Do the necessary changes to the codebase based on the users input.\n" +
        "Don't ask any follow up questions.";
    }

    const escapedPrompt = this.escapePrompt(prompt);
    let _prompt = `${instruction}\n\nUser: ${escapedPrompt}`;

    return {
      command: `echo "${this.escapePrompt(_prompt)}" | gemini --model ${
        this.model || "gemini-2.5-pro"
      } --yolo`,
      errorPrefix: "Gemini",
      labelName: "gemini",
      labelColor: "4285F4",
      labelDescription: "Generated by Gemini AI agent",
    };
  }

  protected getDefaultTemplate(): string {
    return "vibekit-gemini";
  }

  protected getEnvironmentVariables(): Record<string, string> {
    return {
      GEMINI_API_KEY: this.geminiApiKey,
      GOOGLE_API_KEY: this.geminiApiKey, // Some tools might expect this name
    };
  }

  protected getApiKey(): string {
    return this.geminiApiKey;
  }

  protected getAgentType(): AgentType {
    return "gemini";
  }

  protected getModelConfig(): ModelConfig {
    return {
      provider: "gemini",
      apiKey: this.geminiApiKey,
      model: this.model,
    };
  }

  // Override generateCode to support history in the instruction
  public async generateCode(
    prompt: string,
    mode?: "ask" | "code",
    branch?: string,
    history?: Conversation[],
    callbacks?: GeminiStreamCallbacks,
    background?: boolean
  ): Promise<GeminiResponse> {
    let instruction: string;
    if (mode === "ask") {
      instruction =
        "Research the repository and answer the user's questions. " +
        "Do NOT make any changes to any files in the repository.";
    } else {
      instruction =
        "Do the necessary changes to the codebase based on the users input.\n" +
        "Don't ask any follow up questions.";
    }

    if (history && history.length > 0) {
      instruction += `\n\nConversation history: ${history
        .map((h) => `${h.role}\n ${this.escapePrompt(h.content)}`)  // Escape history content too!
        .join("\n\n")}`;
    }

    const escapedPrompt = this.escapePrompt(prompt);

    let _prompt = `${instruction}\n\nUser: ${escapedPrompt}`;

    // Override the command config with history-aware instruction
    const originalGetCommandConfig = this.getCommandConfig.bind(this);
    this.getCommandConfig = (p: string, m?: "ask" | "code") => ({
      ...originalGetCommandConfig(p, m),
      command: `echo "${this.escapePrompt(_prompt)}" | gemini --model ${
        this.model || "gemini-2.5-pro"
      } --yolo`,
    });

    const result = await super.generateCode(
      prompt,
      mode,
      branch,
      history,
      callbacks,
      background
    );

    // Restore original method
    this.getCommandConfig = originalGetCommandConfig;

    return result as GeminiResponse;
  }
}
